home *** CD-ROM | disk | FTP | other *** search
/ Chip 2006 June / CHIP 2006-06.2.iso / program / freeware / Democracy-0.8.2.exe / xulrunner / python / BitTorrent / Storage.py < prev    next >
Encoding:
Python Source  |  2006-04-10  |  10.6 KB  |  286 lines

  1. # The contents of this file are subject to the BitTorrent Open Source License
  2. # Version 1.0 (the License).  You may not copy or use this file, in either
  3. # source code or executable form, except in compliance with the License.  You
  4. # may obtain a copy of the License at http://www.bittorrent.com/license/.
  5. #
  6. # Software distributed under the License is distributed on an AS IS basis,
  7. # WITHOUT WARRANTY OF ANY KIND, either express or implied.  See the License
  8. # for the specific language governing rights and limitations under the
  9. # License.
  10.  
  11. # Written by Bram Cohen
  12.  
  13. import os
  14. from bisect import bisect_right
  15. from array import array
  16.  
  17. from BitTorrent.obsoletepythonsupport import *
  18.  
  19. from BitTorrent import BTFailure
  20.  
  21.  
  22. class FilePool(object):
  23.  
  24.     def __init__(self, max_files_open):
  25.         self.max_files_open = max_files_open
  26.         self.allfiles = {}
  27.         self.handlebuffer = None
  28.         self.handles = {}
  29.         self.whandles = {}
  30.  
  31.     def close_all(self):
  32.         failures = {}
  33.         for filename, handle in self.handles.iteritems():
  34.             try:
  35.                 handle.close()
  36.             except Exception, e:
  37.                 failures[self.allfiles[filename]] = e
  38.         self.handles.clear()
  39.         self.whandles.clear()
  40.         if self.handlebuffer is not None:
  41.             del self.handlebuffer[:]
  42.         for torrent, e in failures.iteritems():
  43.             torrent.got_exception(e)
  44.  
  45.     def set_max_files_open(self, max_files_open):
  46.         self.max_files_open = max_files_open
  47.         self.close_all()
  48.         if len(self.allfiles) > self.max_files_open:
  49.             self.handlebuffer = []
  50.         else:
  51.             self.handlebuffer = None
  52.  
  53.     def add_files(self, files, torrent):
  54.         for filename in files:
  55.             if filename in self.allfiles:
  56.                 raise BTFailure('File '+filename+' belongs to another running '
  57.                                 'torrent')
  58.         for filename in files:
  59.             self.allfiles[filename] = torrent
  60.         if self.handlebuffer is None and \
  61.                len(self.allfiles) > self.max_files_open:
  62.             self.handlebuffer = self.handles.keys()
  63.  
  64.     def remove_files(self, files):
  65.         for filename in files:
  66.             del self.allfiles[filename]
  67.         if self.handlebuffer is not None and \
  68.                len(self.allfiles) <= self.max_files_open:
  69.             self.handlebuffer = None
  70.  
  71.  
  72. # Make this a separate function because having this code in Storage.__init__()
  73. # would make python print a SyntaxWarning (uses builtin 'file' before 'global')
  74.  
  75. def bad_libc_workaround():
  76.     global file
  77.     def file(name, mode = 'r', buffering = None):
  78.         return open(name, mode)
  79.  
  80. class Storage(object):
  81.  
  82.     def __init__(self, config, filepool, files, check_only=False):
  83.         self.filepool = filepool
  84.         self.config = config
  85.         self.ranges = []
  86.         self.myfiles = {}
  87.         self.tops = {}
  88.         self.undownloaded = {}
  89.         self.unallocated = {}
  90.         total = 0
  91.         for filename, length in files:
  92.         # Was self.unallocated[filename] = 0
  93.         # Changed to deal with pre-allocated space
  94.             self.unallocated[filename] = 0
  95.             self.undownloaded[filename] = length
  96.             if length > 0:
  97.                 self.ranges.append((total, total + length, filename))
  98.                 self.myfiles[filename] = None
  99.             total += length
  100.             if os.path.exists(filename):
  101.                 if not os.path.isfile(filename):
  102.                     raise BTFailure('File '+filename+' already exists, but '
  103.                                     'is not a regular file')
  104.                 l = os.path.getsize(filename)
  105.                 if l > length and not check_only:
  106.                     h = file(filename, 'rb+')
  107.                     h.truncate(length)
  108.                     h.close()
  109.                     l = length
  110.         # Changed to pre-allocate space
  111.         if l < length:
  112.             h = file(filename, 'rb+')
  113.             h.seek(length-1)
  114.             h.write('.')
  115.             h.close()
  116.                 self.tops[filename] = length
  117.             elif not check_only:
  118.                 f = os.path.split(filename)[0]
  119.                 if f != '' and not os.path.exists(f):
  120.                     os.makedirs(f)
  121.         # Changed to pre-allocate space
  122.         h = file(filename, 'wb')
  123.         h.seek(length-1)
  124.         h.write('.')
  125.         h.close()
  126.         self.begins = [i[0] for i in self.ranges]
  127.         self.total_length = total
  128.         if check_only:
  129.             return
  130.         self.handles = filepool.handles
  131.         self.whandles = filepool.whandles
  132.  
  133.         # Rather implement this as an ugly hack here than change all the
  134.         # individual calls. Affects all torrent instances using this module.
  135.         if config['enable_bad_libc_workaround']:
  136.             bad_libc_workaround()
  137.  
  138.     def was_preallocated(self, pos, length):
  139.         for filename, begin, end in self._intervals(pos, length):
  140.             if self.tops.get(filename, 0) < end:
  141.                 return False
  142.         return True
  143.  
  144.     def get_total_length(self):
  145.         return self.total_length
  146.  
  147.     def _intervals(self, pos, amount):
  148.         r = []
  149.         stop = pos + amount
  150.         p = bisect_right(self.begins, pos) - 1
  151.         while p < len(self.ranges) and self.ranges[p][0] < stop:
  152.             begin, end, filename = self.ranges[p]
  153.             r.append((filename, max(pos, begin) - begin, min(end, stop) - begin))
  154.             p += 1
  155.         return r
  156.  
  157.     def _get_file_handle(self, filename, for_write):
  158.         handlebuffer = self.filepool.handlebuffer
  159.         if filename in self.handles:
  160.             if for_write and filename not in self.whandles:
  161.                 self.handles[filename].close()
  162.                 self.handles[filename] = file(filename, 'rb+', 0)
  163.                 self.whandles[filename] = None
  164.             if handlebuffer is not None and handlebuffer[-1] != filename:
  165.                 handlebuffer.remove(filename)
  166.                 handlebuffer.append(filename)
  167.         else:
  168.             if for_write:
  169.                 self.handles[filename] = file(filename, 'rb+', 0)
  170.                 self.whandles[filename] = None
  171.             else:
  172.                 self.handles[filename] = file(filename, 'rb', 0)
  173.             if handlebuffer is not None:
  174.                 if len(handlebuffer) >= self.filepool.max_files_open:
  175.                     oldfile = handlebuffer.pop(0)
  176.                     if oldfile in self.whandles:   # .pop() in python 2.3
  177.                         del self.whandles[oldfile]
  178.                     self.handles[oldfile].close()
  179.                     del self.handles[oldfile]
  180.                 handlebuffer.append(filename)
  181.         return self.handles[filename]
  182.  
  183.     def read(self, pos, amount):
  184.         r = []
  185.         for filename, pos, end in self._intervals(pos, amount):
  186.             h = self._get_file_handle(filename, False)
  187.             h.seek(pos)
  188.             r.append(h.read(end - pos))
  189.         r = ''.join(r)
  190.         if len(r) != amount:
  191.             raise BTFailure('Short read - something truncated files?')
  192.         return r
  193.  
  194.     def write(self, pos, s):
  195.         # might raise an IOError
  196.         total = 0
  197.         for filename, begin, end in self._intervals(pos, len(s)):
  198.             h = self._get_file_handle(filename, True)
  199.             h.seek(begin)
  200.             h.write(s[total: total + end - begin])
  201.             total += end - begin
  202.  
  203.     def close(self):
  204.         error = None
  205.         for filename in self.handles.keys():
  206.             if filename in self.myfiles:
  207.                 try:
  208.                     self.handles[filename].close()
  209.                 except Exception, e:
  210.                     error = e
  211.                 del self.handles[filename]
  212.                 if filename in self.whandles:
  213.                     del self.whandles[filename]
  214.         handlebuffer = self.filepool.handlebuffer
  215.         if handlebuffer is not None:
  216.             handlebuffer = [f for f in handlebuffer if f not in self.myfiles]
  217.             self.filepool.handlebuffer = handlebuffer
  218.         if error is not None:
  219.             raise error
  220.  
  221.     def write_fastresume(self, resumefile, amount_done):
  222.         resumefile.write('BitTorrent resume state file, version 1\n')
  223.         resumefile.write(str(amount_done) + '\n')
  224.         for _, _, filename in self.ranges:
  225.             resumefile.write(str(os.path.getsize(filename)) + ' ' +
  226.                              str(os.path.getmtime(filename)) + '\n')
  227.  
  228.     def check_fastresume(self, resumefile, return_filelist=False,
  229.                          piece_size=None, numpieces=None, allfiles=None):
  230.         filenames = [name for _, _, name in self.ranges]
  231.         if resumefile is not None:
  232.             version = resumefile.readline()
  233.             if version != 'BitTorrent resume state file, version 1\n':
  234.                 raise BTFailure('Unsupported fastresume file format, '
  235.                       'maybe from another client version')
  236.             amount_done = int(resumefile.readline())
  237.         else:
  238.             amount_done = size = mtime = 0
  239.         for filename in filenames:
  240.             if resumefile is not None:
  241.                 line = resumefile.readline()
  242.                 size, mtime = line.split()[:2] # allow adding extra fields
  243.                 size = int(size)
  244.                 mtime = int(mtime)
  245.             if os.path.exists(filename):
  246.                 fsize = os.path.getsize(filename)
  247.             else:
  248.                 fsize = 0
  249.             if fsize > 0 and mtime != os.path.getmtime(filename):
  250.                 raise BTFailure("Fastresume info doesn't match file "
  251.                                 "modification time")
  252.             if size != fsize:
  253.                 raise BTFailure("Fastresume data doesn't match actual "
  254.                                 "filesize")
  255.         if not return_filelist:
  256.             return amount_done
  257.         if resumefile is None:
  258.             return None
  259.         if numpieces < 32768:
  260.             typecode = 'h'
  261.         else:
  262.             typecode = 'l'
  263.         try:
  264.             r = array(typecode)
  265.             r.fromfile(resumefile, numpieces)
  266.         except Exception, e:
  267.             raise BTFailure("Couldn't read fastresume data: " + str(e))
  268.         for i in range(numpieces):
  269.             if r[i] >= 0:
  270.                 # last piece goes "past the end", doesn't matter
  271.                 self.downloaded(r[i] * piece_size, piece_size)
  272.             if r[i] != -2:
  273.                 self.allocated(i * piece_size, piece_size)
  274.         undl = self.undownloaded
  275.         unal = self.unallocated
  276.         return amount_done, [undl[x] for x in allfiles], \
  277.                [not unal[x] for x in allfiles]
  278.  
  279.     def allocated(self, pos, length):
  280.         for filename, begin, end in self._intervals(pos, length):
  281.             self.unallocated[filename] -= end - begin
  282.  
  283.     def downloaded(self, pos, length):
  284.         for filename, begin, end in self._intervals(pos, length):
  285.             self.undownloaded[filename] -= end - begin
  286.